import pickle

import numpy as np
import scipy.io.wavfile as wf
import python_speech_features as sf
import sklearn.svm as svm
import sklearn.metrics as sm
import sklearn.preprocessing as sp
import os
import librosa.display


# 整理样本
def search_files(directory):
    """
        检索目录下的所有wav文件 返回目录字典
        {“appple”:[url, url...],
        “kiwi”:[url, url...],....}
    """
    files_dict = {}
    for cur_dir, sub_dirs, files in os.walk(directory):  #cur_dir：当前目录，sub_dirs：下一级目录，file:音频文件
        for file in files:
            if file.endswith(".wav"):
                label = cur_dir.split(os.path.sep)[-1]
                if label not in files_dict:
                    files_dict[label] = []
                url = os.path.join(cur_dir, file)  #音频文件路径
                files_dict[label].append(url)   #files_dict：音频文件路径与label一一对应
    return files_dict


def files_mfc(file_urls):
    """
        读取语音文件 并且 MFCC化音频文件 梅尔频率倒谱系数
        返回整理好的 样本数据
    """
    x_data, y_data = [], []
    for label, urls in file_urls.items():  # 音频文件和label组成一个元组
        for file in urls:
            sample_rate, signs = wf.read(file)
            mfc = sf.mfcc(signs, sample_rate)
            #print(sample_rate)
            # ==>(1,13)
            sample = np.mean(mfc, axis=0)  # 求平均值
            x_data.append(sample)   # 音频特征数组
            y_data.append(label)    # 标签特征数组
    x_data = np.array(x_data)
    return x_data, y_data


# 读取训练集 文件路径
train_urls = search_files("D:/Pycharm/trainvoice_depoyment/train")   #train_urls<——files_dict，音频文件路径与label一一对应
print(train_urls)

# 整理训练集
train_x, train_y = files_mfc(train_urls)   #train_x, train_y为梅尔倒频谱处理后的训练数据，x为数据，y为标签
#print(train_x)
# 标签编码
encoder = sp.LabelEncoder()
train_y_label = encoder.fit_transform(train_y)  #处理标签
print(train_x.shape, train_y_label.shape)

# 创建模型 SVM 并训练
model = svm.SVC(kernel="poly", degree=2,C=2, gamma="auto", probability=True)
model.fit(train_x, train_y_label)
file = open(r"D:\graduation design\model_deployment\voc_svm\voc_svm\model_mfcc.pickle", "wb")
pickle.dump(model, file)
file.close()

# # 读取测试集 文件路径
# test_urls = search_files("D:/Pycharm/trainvoice_max/test")
# print(test_urls)
#
# # 整理测试集
# test_x, test_y = files_mfc(test_urls)
# test_y_label = encoder.fit_transform(test_y)
#
# # 预测测试集 输出分类报告
# prd_test_y = model.predict(test_x)
# print(sm.classification_report(test_y_label, prd_test_y))

# # 输出置信概率
# probs = model.predict_proba(test_x)
# print(np.round(probs, 3))  #概率probs保留3位小数
#
# # 循环打印结果
# for label, prob in zip(encoder.inverse_transform(prd_test_y), probs.max(axis=1)):  #test_y——>test_y_label——>test_y
#     print(label, np.round(prob, 3))

#
# #绘制热力混淆矩阵
# import numpy as np
# import seaborn as sns
# import matplotlib.pyplot as plt
# from sklearn.metrics import confusion_matrix
# # 创建标签映射字典
# label_map = {label: i for i, label in enumerate(encoder.classes_)}
#
# # 将预测和真实标签转换为新的标签
# prd_test_y_mapped = [label_map[label] for label in encoder.inverse_transform(prd_test_y)]
# test_y_label_mapped = [label_map[label] for label in encoder.inverse_transform(test_y_label)]
#
# # 计算混淆矩阵
# conf_mat = confusion_matrix(test_y_label_mapped, prd_test_y_mapped)
#
# # 绘制热力混淆矩阵
# plt.figure(figsize=(10, 8))
# sns.set(font_scale=1.2)  # 设置字体大小
# sns.heatmap(conf_mat, annot=True, fmt='d', cmap="YlGnBu",
#             xticklabels=np.arange(7), yticklabels=np.arange(7))
# plt.xlabel('Predicted labels')
# plt.ylabel('True labels')
# plt.title('Confusion Matrix')
# plt.show()